import numpy as np 

class MyHMM():

    def __init__(self, hmm):
        self.hmm = hmm

    # 给定一个观测值 返回它从所有样本中 可能的概率分布
    def prob_O2S(self, model, o):
        M_O2S = self.hmm["M_O2S"]
        return M_O2S[:,int(o)]

    def gen_samples_from_HMM(self, N):
        M_O2S = self.hmm["M_O2S"]
        
        datas = np.zeros(N)
        status = np.zeros(N)

        #得到初始状态 并根据初始状态生成一个样本
        init_S = self.gen_one_sample_from_Prob_list(self.hmm["pi"])
        status[0] = init_S

        #根据初始状态 生成一个数据
        datas[0] = self.gen_one_sample_from_Prob_list(M_O2S[int(status[0])])

        # 生成其他样本
        for i in range(1,N):
            # 根据前一个状态，生成当前的状态
            status[i] = self.gen_one_sample_from_Prob_list(self.hmm["A"][int(status[i-1])])
            # 根据当前的状态生成一个数据
            datas[i] = self.gen_one_sample_from_Prob_list(M_O2S[int(status[i])])
        return datas, status

    # 从一个给定的概率列表中 给出一个符合其概率分布的样本
    def gen_one_sample_from_Prob_list(self, prob_list):
        N_segment = np.shape(prob_list)[0]  
        prob_segment = np.zeros(N_segment)
        for i in range(N_segment):
            prob_segment[i] = prob_segment[i-1] + prob_list[i]
        s = 0
        for i in range(N_segment):
            data = np.random.rand()
            if data <= prob_segment[i]:
                s = i
                break
        return s


    # 用来事后判定是否与状态的转移矩阵符合
    def judge_consistent_with_state_transition_matrix(self, S, N):
        datas, status = self.gen_samples_from_HMM(int(N))

        # 进行概率的统计
        s = int(S)
        index_current_s = np.where(status==s)[0]
        index_next_s = index_current_s + 1
        index_next_s = index_next_s[:-1]

        temp = status[index_next_s]
        print("Status %d to"%(s))
        print("%.3f"%(np.where(temp==0)[0].shape[0]/np.shape(temp)[0]))
        print("%.3f"%(np.where(temp==1)[0].shape[0]/np.shape(temp)[0]))
        print("%.3f"%(np.where(temp==2)[0].shape[0]/np.shape(temp)[0]))

        # 对不同状态下的数据分布进行统计
        temp = datas[index_current_s]
        print("datas at status %d"%(s))
        print("%.3f"%(np.where(temp==0)[0].shape[0]/np.shape(temp)[0]))
        print("%.3f"%(np.where(temp==1)[0].shape[0]/np.shape(temp)[0]))
        print("%.3f"%(np.where(temp==2)[0].shape[0]/np.shape(temp)[0]))
        print("%.3f"%(np.where(temp==3)[0].shape[0]/np.shape(temp)[0]))
        print("%.3f"%(np.where(temp==4)[0].shape[0]/np.shape(temp)[0]))
        print("%.3f"%(np.where(temp==5)[0].shape[0]/np.shape(temp)[0]))
        print("%.3f"%(np.where(temp==6)[0].shape[0]/np.shape(temp)[0]))
        print("%.3f"%(np.where(temp==7)[0].shape[0]/np.shape(temp)[0]))


if __name__ == '__main__':
    model_hmm = dict()
    # 各个状态初始分布
    model_hmm["pi"] = np.array([1.0/3.0, 1.0/3.0, 1.0/3.0])
    # 行表示当前的状态 列表示下一个状态
    model_hmm["A"] = np.array([
                        [0.4, 0.3, 0.3],
                        [0.3, 0.4, 0.3],
                        [0.4, 0.3, 0.4],
                    ])
    # 观测样本与各个状态之间的概率映射
    M_O2S = np.zeros([3,8])
    M_O2S[0,:4] = 1/4.0
    M_O2S[1,:6] = 1/6.0
    M_O2S[2,:8] = 1/8.0
    model_hmm["M_O2S"] = M_O2S
    hmm = MyHMM(hmm=model_hmm)

    # print(hmm.gen_samples_from_HMM(10))
    hmm.judge_consistent_with_state_transition_matrix(1, 2000)